import numpy as np
from scipy.interpolate import RegularGridInterpolator

from GenericHdf import GenericHdf
from Timer import Timer
from utils import find_files

GLOBAL_LATLON_AREA = (-180, -90, 180, 90)
# 河南省经纬度范围(110, 31, 117, 37)
HN_LATLON_AREA = (110, 31, 117, 37)


class AreaEarthObData:

    def __init__(self, latlon_area=HN_LATLON_AREA):
        self.earth_ob_data = None
        self.latlon_area = latlon_area  # (min_lon, min_lat, max_lon, max_lat)

    def in_area(self, lat, lon):
        min_lon, min_lat, max_lon, max_lat = self.latlon_area
        return min_lon <= lon < max_lon and min_lat <= lat < max_lat

    # 将目标经纬度转为行列坐标，返回浮点数(x, y)，由小数部分表示匹配偏差（后续插值用）
    def latlon_to_xy(self, lat, lon):
        # 确保数据已经存在
        if self.earth_ob_data is None:
            raise ValueError("earth_ob_data is not initialized")
        min_lon, min_lat, max_lon, max_lat = self.latlon_area
        layers, cols, rows = self.earth_ob_data.shape
        lon_span = max_lon - min_lon
        lat_span = max_lat - min_lat
        # 计算x（经度方向，对应列数）
        x = (lon - min_lon) / lon_span * cols
        # 计算y（纬度方向，对应行数）
        y = (max_lat - lat) / lat_span * rows
        return x, y

    # 求行列坐标的观测数据，行列坐标为浮点数(x, y)，根据坐标值的小数部分对每一层的数据进行插值运算
    def earth_ob_at_xy(self, x, y):
        layers, cols, rows = self.earth_ob_data.shape
        results = []
        for i in range(layers):
            interpolator = RegularGridInterpolator((np.arange(cols), np.arange(rows)), self.earth_ob_data[i])
            # interpolator = PartialValidNeighborInterpolator(self.earth_ob_data[i])
            result = interpolator((x, y))
            results.append(result)
        return np.array(results)

    def earth_ob_at(self, lat, lon):
        x, y = self.latlon_to_xy(lat, lon)
        return self.earth_ob_at_xy(x, y)


class Fy3dL1Hdf(GenericHdf, AreaEarthObData):

    def __init__(self, file_path, latlon_area=HN_LATLON_AREA):
        GenericHdf.__init__(self, file_path)
        AreaEarthObData.__init__(self, latlon_area)
        self.file_path = file_path
        self.EARTH_OBSERVE_BT_10_to_89GHz = self.file["Calibration/EARTH_OBSERVE_BT_10_to_89GHz"]
        self.earth_ob_data = Fy3dL1Hdf.par(self.EARTH_OBSERVE_BT_10_to_89GHz)
        self.Latitude = self.file["Geolocation/Latitude"]
        self.lat_data = Fy3dL1Hdf.par(self.Latitude)
        self.Longitude = self.file["Geolocation/Longitude"]
        self.lon_data = Fy3dL1Hdf.par(self.Longitude)

    @classmethod
    def par(cls, dataset):
        fill_value = dataset.attrs.get("FillValue", -32767)
        intercept = dataset.attrs.get("Intercept", 327.68)
        slope = dataset.attrs.get("Slope", 0.01)
        data = np.array(dataset).astype(np.float32)
        indices = np.where(data == fill_value)
        data[indices] = np.nan
        return data * slope + intercept

    def map_to_area(self, latlon_area=None, earth_ob_data_or_shape=None, merge_func=None):
        if latlon_area is None:
            latlon_area = self.latlon_area

        # 解析经纬范围参数
        min_lon, min_lat, max_lon, max_lat = latlon_area
        lon_span = max_lon - min_lon
        lat_span = max_lat - min_lat

        # 确定目标数据容器
        layers = self.earth_ob_data.shape[0]
        if earth_ob_data_or_shape is None:
            # 自动计算网格尺寸（经度跨度*100, 纬度跨度*100）
            cols = int(lon_span * 100)
            rows = int(lat_span * 100)
            earth_ob_data = np.full((layers, cols, rows), np.nan, self.earth_ob_data.dtype)
        elif isinstance(earth_ob_data_or_shape, tuple):
            # 使用指定的形状创建容器
            cols, rows = earth_ob_data_or_shape
            earth_ob_data = np.full((layers, cols, rows), np.nan, self.earth_ob_data.dtype)
        else:
            # 直接使用现有容器
            earth_ob_data = earth_ob_data_or_shape
            cols = earth_ob_data.shape[1]
            rows = earth_ob_data.shape[2]

        # 设置默认合并函数（加权平均）
        if merge_func is None:
            def merge_func(existing_value, new_value):
                if np.isnan(existing_value):
                    return new_value
                if np.isnan(new_value):
                    return existing_value
                return (existing_value + new_value) / 2

        # 遍历所有数据点，lat_data和lon_data的形状一致
        for i in range(self.lat_data.shape[0]):
            for j in range(self.lat_data.shape[1]):
                lat = self.lat_data[i, j]
                lon = self.lon_data[i, j]

                # 跳过无效坐标
                if np.isnan(lat) or np.isnan(lon):
                    continue

                # 检查坐标是否在目标区域内
                if not (min_lon <= lon < max_lon and min_lat <= lat < max_lat):
                    continue

                # 计算网格坐标（考虑浮点精度和边界）
                x = int(((lon - min_lon) / lon_span) * cols)
                y = int(((max_lat - lat) / lat_span) * rows)

                # 执行数据映射
                for k in range(layers):
                    existing = earth_ob_data[k, x, y]
                    new = self.earth_ob_data[k, i, j]
                    earth_ob_data[k, x, y] = merge_func(existing, new)

        # 维护全局数据状态（与原逻辑保持一致）
        if self.earth_ob_data is None:
            self.earth_ob_data = earth_ob_data

        return earth_ob_data

    def map_to_global(self, earth_ob_data_or_shape=None, merge_func=None):
        return self.map_to_area((-180, 90, 180, -90), earth_ob_data_or_shape, merge_func=merge_func)


class Fy3dL1HdfBatch(AreaEarthObData):
    def __init__(self, fy3d_l1_dir, latlon_area=HN_LATLON_AREA):
        AreaEarthObData.__init__(self, latlon_area)
        file_paths = find_files(fy3d_l1_dir, lambda f: '_L1_' in f.upper() and f.upper().endswith('.HDF'))
        self.fy3d_l1_hdfs = []
        for file_path in file_paths:
            self.fy3d_l1_hdfs.append(Fy3dL1Hdf(file_path, latlon_area))
        self.verbose = True
        pass

    def map_to_area(self, latlon_area=None, earth_ob_data_or_shape=None, merge_func=None):
        if self.fy3d_l1_hdfs is None or len(self.fy3d_l1_hdfs) == 0:
            return None
        sample = self.fy3d_l1_hdfs[0]

        if self.verbose:
            timer = Timer()
            timer.tick("map_to_area 开始")

        if latlon_area is None:
            latlon_area = self.latlon_area

        # 解析经纬范围参数
        min_lon, min_lat, max_lon, max_lat = latlon_area
        lon_span = max_lon - min_lon
        lat_span = max_lat - min_lat

        # 确定目标数据容器
        layers = sample.earth_ob_data.shape[0]
        if earth_ob_data_or_shape is None:
            # 自动计算网格尺寸（经度跨度*100, 纬度跨度*100）
            cols = int(lon_span * 100)
            rows = int(lat_span * 100)
            earth_ob_data = np.full((layers, cols, rows), np.nan, sample.earth_ob_data.dtype)
        elif isinstance(earth_ob_data_or_shape, tuple):
            # 使用指定的形状创建容器
            cols, rows = earth_ob_data_or_shape
            earth_ob_data = np.full((layers, cols, rows), np.nan, sample.earth_ob_data.dtype)
        else:
            # 直接使用现有容器
            earth_ob_data = earth_ob_data_or_shape

        if merge_func is None:
            def merge_func(existing_value, index, new_value):
                if np.isnan(existing_value):
                    return new_value
                if np.isnan(new_value):
                    return existing_value
                return (existing_value * index + new_value) / (index + 1)

        for index, fy3d_l1_hdf in enumerate(self.fy3d_l1_hdfs):
            if self.verbose:
                timer.tick(f"开始映射 '{fy3d_l1_hdf.file_path}'")
            fy3d_l1_hdf.map_to_area(latlon_area, earth_ob_data,
                                    lambda existing_value, new_value: merge_func(existing_value, index, new_value))
            if self.verbose:
                timer.tick(f"'{fy3d_l1_hdf.file_path}' 完成")

        if self.earth_ob_data is None:
            self.earth_ob_data = earth_ob_data
        if self.verbose:
            timer.tick("map_to_global 完成")
        return earth_ob_data
